from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.utils import new_agent_text_message


class WritingAgent:
    """
    智能写作 Agent
    """

    async def invoke(self, query=None) -> str:
        """
        智能体处理逻辑
        :return:
        """
        return f"写作智能体已处理：{query}"


class WritingAgentExecutor(AgentExecutor):
    """
    测试 AgentProxy 实现
    """

    def __init__(self):
        self.agent = WritingAgent()

    async def execute(
            self,
            context: RequestContext,
            event_queue: EventQueue,
    ) -> None:
        part_list = context.message.parts
        if part_list:
            text_list = list()
            for part in part_list:
                text_list.append(part.root.text)

            result = await self.agent.invoke(text_list)
        else:
            result = await self.agent.invoke()
        await event_queue.enqueue_event(new_agent_text_message(result))

    async def cancel(
            self, context: RequestContext, event_queue: EventQueue
    ) -> None:
        raise Exception('cancel not supported')
